/*
 * Copyright 2010-2012 Susanta Tewari. <freecode4susant@users.sourceforge.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package coalescent.is.cef.command.listener.exe;

import bd.org.apache.commons.math.stat.descriptive.SummaryStatistics;
import cef.command.Command;
import cef.command.listener.exe.CommandExeEvent;
import cef.command.listener.exe.FigureEditorExeComponent;
import coalescent.is.cef.command.RunSamplers_Cmd;
import coalescent.is.cef.command.RunSamplers_Multi_Cmd;
import coalescent.is.cef.event.RunSamplers_Multi_Evt;
import commons.util.ApfloatUtil;
import commons.util.ArrayUtil;
import org.apfloat.Apcomplex;
import org.apfloat.Apfloat;
import org.apfloat.ApfloatMath;
import org.jfree.chart.ChartPanel;
import org.jfree.chart.JFreeChart;
import org.jfree.chart.axis.NumberAxis;
import org.jfree.chart.plot.XYPlot;
import org.jfree.chart.renderer.xy.XYLineAndShapeRenderer;
import org.jfree.data.general.DatasetChangeEvent;
import org.jfree.data.xy.DefaultXYDataset;

import javax.swing.*;
import java.awt.*;
import java.math.BigDecimal;
import java.text.DecimalFormat;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.Map;

import static commons.util.ApfloatUtil.Max;
import static commons.util.ApfloatUtil.Min;
import static commons.util.FormatUtil.format;

/**
 * @since
 * @version
 * @author Susanta Tewari
 * @history Created on 2/4/13.
 */
public abstract class VizSamplers_Multi_EC extends FigureEditorExeComponent {

    /** formatting */
    private static final Apfloat LOG_BASE  = new Apfloat(10);
    private static final DecimalFormat fmt = new DecimalFormat("0.####");

    /** event cache */
    protected Collection<String> samplers;
    private double[] params;

    /** sampler maps & data set */
    private Map<String, JLabel> itr_map;
    private Map<String, JLabel> measure_map;
    private Map<String, JLabel> mme_map;

    /** GUI settings */
    protected RunSamplers_Multi_Evt.GUI guiSettings;
    protected Map<String, Color> colorMap;
    private static final Stroke stroke = new BasicStroke(4.0f);

    /** chart */
    private Map<String, double[][]> data_map;
    protected DefaultXYDataset data_set;
    private Apfloat rangeMin = Apfloat.ZERO,
                    rangeMax = Apfloat.ZERO;    // dynamically adjusted
    private JFreeChart chart;

    /** update-1 controls */
    protected final JPanel update1_panel      = new JPanel();
    private static final GridBagConstraints c = new GridBagConstraints();

    /** update-2 controls */
    private int timer2Counter;

    protected void updatedRangeMin(final Apfloat newVal) {
        rangeMin = ApfloatUtil.Min(rangeMin, newVal);
    }

    protected void updatedRangeMax(final Apfloat newVal) {
        rangeMax = ApfloatUtil.Max(rangeMax, newVal);
    }

    @Override
    public void initComponents(final Command cmd) {

        final RunSamplers_Multi_Cmd command = (RunSamplers_Multi_Cmd) cmd;

        guiSettings = command.getEvent().getGUISettings();


        // init params
        final BigDecimal[] params1 = command.getEvent().getParams();

        params = new double[params1.length];

        for (int i = 0; i < params1.length; i++) {
            params[i] = params1[i].doubleValue();
        }


        // sampler maps & data set
        colorMap    = guiSettings.getColorMap();
        samplers    = command.getSamplerNames();
        itr_map     = new HashMap<>(samplers.size());
        mme_map     = new HashMap<>(samplers.size());
        measure_map = new HashMap<>(samplers.size());
        data_map    = new HashMap<>(samplers.size());
        data_set    = new DefaultXYDataset();

        for (final String samplerName : samplers) {

            final JLabel itr_jLabel = of_label();
            final JLabel mme_jLabel = of_label();

            itr_jLabel.setForeground(colorMap.get(samplerName));
            mme_jLabel.setForeground(colorMap.get(samplerName));
            itr_map.put(samplerName, itr_jLabel);
            mme_map.put(samplerName, mme_jLabel);

            final double[][] data = create_data(params.length);

            data_map.put(samplerName, data);
            data_set.addSeries(samplerName, data);
        }


        // init chart
        title = getTitle();

        final NumberAxis x_axis = new NumberAxis("Theta");
        final NumberAxis y_axis = new NumberAxis("Probability");

        x_axis.setAutoRangeIncludesZero(false);
        y_axis.setAutoRangeIncludesZero(false);

        final XYLineAndShapeRenderer renderer = new XYLineAndShapeRenderer(true, false);
        final XYPlot xyplot                   = new XYPlot(data_set, x_axis, y_axis, renderer);

        for (final String s : samplers) {

            final int series = data_set.indexOf(s);

            renderer.setSeriesPaint(series, colorMap.get(s));
            renderer.setSeriesStroke(series, stroke);
        }

        chart  = new JFreeChart(title, JFreeChart.DEFAULT_TITLE_FONT, xyplot, true);
        figure = new JPanel(new BorderLayout());

        figure.add(new ChartPanel(chart), BorderLayout.CENTER);


        // timer1 panel
        update1_panel.setLayout(new GridBagLayout());

        c.insets = new Insets(5, 5, 5, 5);
        c.gridx  = 0;
        c.gridy  = 0;


        // iterations
        update1_panel.add(new JLabel("Iterations: "), c);

        for (final String s : samplers) {

            c.gridy++;

            update1_panel.add(itr_map.get(s), c);
        }


        // mme
        c.gridx = 1;
        c.gridy = 0;

        final JLabel mme_label = new JLabel("MME: ");

        mme_label.setToolTipText("Maximum Mean Estimator");
        update1_panel.add(mme_label, c);

        for (final String s : samplers) {

            c.gridy++;

            update1_panel.add(mme_map.get(s), c);
        }


        // measure
        c.gridx = 2;
        c.gridy = 0;

        update1_panel.add(new JLabel(getMeasureTitle() + ": "), c);

        for (final String s : samplers) {

            c.gridy++;

            final JLabel jLabel = of_label();

            jLabel.setForeground(colorMap.get(s));
            measure_map.put(s, jLabel);
            update1_panel.add(jLabel, c);
        }

        figure.add(update1_panel, BorderLayout.EAST);
    }

    @Override
    public void updateComponentData(final CommandExeEvent event) {

        final RunSamplers_Multi_Cmd command = (RunSamplers_Multi_Cmd) event.getSource();

        update_1(command);

        if (timer2Counter % 3 == 0) update_2(command);

        timer2Counter = (timer2Counter >= 100) ? 0 : timer2Counter + 1;
    }

    @Override
    public void lastUpdate(final CommandExeEvent event) {

        final RunSamplers_Multi_Cmd command = (RunSamplers_Multi_Cmd) event.getSource();

        update_1(command);

        for (int i = 0; i < 3; i++) {
            update_2(command);    // plot requires min 3 points to display
        }
    }

    /**
     * Updates iteration count and IS measure.
     *
     * @param cmd command
     */
    protected void update_1(final RunSamplers_Multi_Cmd cmd) {

        for (final String sampler : samplers) {

            final long itrCount = cmd.getSampleSize(sampler);

            if (itrCount < 3) return;

            final JLabel jLabel = itr_map.get(sampler);

            jLabel.setText("<html>" + "<big>[" + format(itrCount) + "]</big></html>");
        }
    }

    /**
     * updates graph.
     *
     * @param cmd command
     */
    protected void update_2(final RunSamplers_Multi_Cmd cmd) {

        final SummaryStatistics stat = new SummaryStatistics();

        rangeMin = Apcomplex.ZERO;
        rangeMax = Apcomplex.ZERO;

        for (final String s : samplers) {

            final Apfloat[] estimate    = getISMeasure(cmd, s);
            final int index_mme         = ArrayUtil.indexOfMax(estimate);
            final double est            = params[index_mme];
            final double v              = estimate[index_mme].doubleValue();
            final JLabel mme_jLabel     = mme_map.get(s);
            final JLabel measure_jLabel = measure_map.get(s);

            mme_jLabel.setText("<html>" + "<big>[" + fmt.format(est) + "]</big></html>");
            measure_jLabel.setText("<html>" + "<big>[" + fmt.format(v) + "]</big></html>");

            if (estimate == null) break;    // cases for log zero etc..


            // add data to data set
            double[][] data = data_map.get(s);

            for (int i = 0; i < estimate.length; i++) {

                data[1][i] = estimate[i].doubleValue();

                stat.addValue(data[1][i]);
            }


            // range min,max
            final Apfloat min = new Apfloat(stat.getMin());
            final Apfloat max = new Apfloat(stat.getMax());

            rangeMin = (rangeMin == Apcomplex.ZERO) ? min : Min(rangeMin, min);
            rangeMax = (rangeMax == Apcomplex.ZERO) ? max : Max(rangeMax, max);

            data_set.addSeries(s, data);
        }


        // adjust range axis
        final double se = stat.getStandardDeviation().doubleValue();

        chart.getXYPlot().getRangeAxis().setRange(rangeMin.doubleValue() - se,
                rangeMax.doubleValue() + se);


        // set the chart data set
        chart.getXYPlot().datasetChanged(new DatasetChangeEvent(this, data_set));
    }

    private double[][] create_data(int PARAM_LENGTH) {

        final double[][] data = new double[2][PARAM_LENGTH];

        for (int i = 0; i < params.length; i++) {
            data[0][i] = params[i];
        }

        Arrays.fill(data[1], Double.NaN);

        return data;
    }

    private JLabel of_label() {

        final JLabel jLabel = new JLabel();

        jLabel.setOpaque(true);
        jLabel.setBackground(Color.white);

        return jLabel;
    }

    /**
     * @return compute the corresponding IS measure
     */
    protected abstract Apfloat[] getISMeasure(RunSamplers_Multi_Cmd command, String samplerName);

    protected abstract String getMeasureTitle();

    protected abstract String getTitle();

    protected abstract String getYAxisTitle();

    @Override
    public Class<? extends Command> getTargetCommandClass() {
        return RunSamplers_Cmd.class;
    }

    /**
     * IS ESS tracker
     */
    public static final class IS_TRACKER_ESS extends VizSamplers_Multi_EC {

        @Override
        protected Apfloat[] getISMeasure(final RunSamplers_Multi_Cmd command,
                                         final String samplerName) {

            final BigDecimal[] ess = command.getIS_ESS(samplerName);
            final Apfloat[] result = new Apfloat[ess.length];

            for (int i = 0; i < ess.length; i++) {
                result[i] = new Apfloat(ess[i]);
            }

            return result;
        }

        @Override
        protected String getMeasureTitle() {
            return "ESS";
        }

        @Override
        protected String getTitle() {
            return "Compare Importance Samplers on ESS";
        }

        @Override
        protected String getYAxisTitle() {
            return "ESS";
        }

        @Override
        public IS_TRACKER_ESS newInstance() {
            return new IS_TRACKER_ESS();
        }
    }

    /**
     * IS Mean tracker
     */
    public static final class IS_TRACKER_MEAN extends VizSamplers_Multi_EC {

        @Override
        protected Apfloat[] getISMeasure(final RunSamplers_Multi_Cmd command,
                                         final String samplerName) {
            return to_Log(command.getISMean(samplerName));
        }

        @Override
        protected String getMeasureTitle() {
            return "Means";
        }

        @Override
        protected String getTitle() {
            return "Compare Importance Samplers on Mean";
        }

        @Override
        protected String getYAxisTitle() {
            return "Lg Mean";
        }

        @Override
        public IS_TRACKER_MEAN newInstance() {
            return new IS_TRACKER_MEAN();
        }
    }

    private static Apfloat[] to_Log(BigDecimal[] mean) {

        final Apfloat[] result = new Apfloat[mean.length];

        for (int i = 0; i < mean.length; i++) {
            result[i] = ApfloatMath.log(new Apfloat(mean[i]), LOG_BASE);
        }

        return result;
    }

    /**
     * IS SE tracker
     */
    public static final class IS_TRACKER_SE extends VizSamplers_Multi_EC {

        @Override
        protected Apfloat[] getISMeasure(final RunSamplers_Multi_Cmd command,
                                         final String samplerName) {
            return to_Log(command.getIS_SE(samplerName));
        }

        @Override
        protected String getMeasureTitle() {
            return "Std. Error";
        }

        @Override
        protected String getTitle() {
            return "Compare Importance Samplers on SE";
        }

        @Override
        protected String getYAxisTitle() {
            return "Lg SE";
        }

        @Override
        public IS_TRACKER_SE newInstance() {
            return new IS_TRACKER_SE();
        }
    }
}
